A Fix for Softmax and Cross-Entropy
softmax 与交叉熵的数值修复：消除梯度塌缩与抵消
🧭 数值稳定训练 · softmax/CE 修复 ★ 深度解读
Google DeepMind
修复 softmax/cross-entropy 的两个数值 bug，消除 Slingshot 尖刺与蒸馏漂移
📌 工作参考价值（4 点）
- 在大规模多分类/生成式推荐训练（含 bf16）中，若出现 loss 精确为 0、梯度不对称或 Slingshot 尖刺，优先检查 logit margin 是否超过精度悬崖：bf16 约 5.5，fp32 约 16.6。可直接替换 CE 为论文的稳定实现：loss=softplus(-margin)，margin=-logsumexp(非目标 logits - 目标 logit)，保证 target logit 梯度不归零且梯度严格零和。
- 做知识蒸馏（如 LLM 教师蒸馏排序模型、语义 ID 蒸馏、ctr 蒸馏）时，学生接近教师时标准 soft-target CE 会因 q-p 相减产生非零噪声梯度，导致学生漂移；应改用 log-space 直接计算 q-p 的实现（expm1/log1p），保证 zS=zT 时梯度为 0，bf16 下尤其重要。
- AdamW/Adam 训练时 logit margin 最终会稳定在 log(1/ε)，若 ε=1e-8 则 margin≈18.4，略高于 fp32 悬崖；可通过调整 ε 避开精度悬崖或使用稳定 CE，避免高置信度下的隐性优化问题。
- 该修复为 drop-in 替换、无额外计算开销，建议直接用于训练框架的 CE 实现，防止高置信度场景的数值退化。